package chainx

import (
	"context"
	"errors"

	"github.com/tmc/langchaingo/callbacks"
	"github.com/tmc/langchaingo/chains"
	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/outputparser"
	"github.com/tmc/langchaingo/schema"
)

const Empty = "DEFAULT"

var ErrNotHandles = errors.New("不存在合适的handler")

type Handler interface {
	Name() string
	Description() string
	Call(ctx context.Context, inputs map[string]any, options ...chains.ChainCallOption) (map[string]any, error)
}

type Router struct {
	Handles      map[string]Handler
	Chain        chains.Chain
	callbacks    callbacks.Handler
	memory       schema.Memory
	OutputParser outputparser.Structured
	emptyHandle  Handler
}

func NewRouter(llm llms.Model, handlers []Handler, opts ...RouterOption) *Router {
	opt := executorDefaultOptions(handlers)
	for _, o := range opts {
		o(&opt)
	}

	hs := make(map[string]Handler, len(handlers))
	for _, handler := range handlers {
		hs[handler.Name()] = handler
	}

	return &Router{
		Handles: hs,
		Chain: chains.NewLLMChain(
			llm,
			opt.prompt,
		),
		callbacks:    opt.callbacksHandler,
		memory:       opt.memory,
		emptyHandle:  opt.emptyHandle,
		OutputParser: _outputParser,
	}
}

func (c *Router) Call(ctx context.Context, inputs map[string]any, options ...chains.ChainCallOption) (map[string]any, error) {

	if c.callbacks != nil {
		c.callbacks.HandleChainStart(ctx, inputs)
	}

	if len(c.Handles) == 0 {
		if c.emptyHandle != nil {
			return c.emptyHandle.Call(ctx, inputs, options...)
			// return chains.Call(ctx, c.emptyHandle.Chains(), inputs)
		} else {
			return nil, ErrNotHandles
		}
	}

	result, err := chains.Call(ctx, c.Chain, inputs, options...)
	if err != nil {
		return nil, err
	}

	text, ok := result["text"]
	if !ok {
		return nil, chains.ErrNotFound
	}

	out, err := c.OutputParser.Parse(text.(string))
	if err != nil {
		return nil, err
	}
	if c.callbacks != nil {
		c.callbacks.HandleChainEnd(ctx, map[string]any{
			"out": out,
		})
	}

	data := out.(map[string]string)
	next, ok := data[_destinations]
	if !ok || next == Empty || c.Handles[next] == nil {
		if c.emptyHandle != nil {
			return c.emptyHandle.Call(ctx, inputs, options...)
			// return chains.Call(ctx, c.emptyHandle.Chains(), inputs)
		} else {
			return nil, ErrNotHandles
		}
	}

	inputs[_input] = data[_input]

	// return chains.Call(ctx, c.Handles[next].Chains(), inputs)
	return c.Handles[next].Call(ctx, inputs, options...)
}

// GetMemory gets the memory of the chain.
func (c *Router) GetMemory() schema.Memory {
	return c.memory
}

// GetInputKeys returns the input keys the chain expects.
func (c *Router) GetInputKeys() []string {
	return []string{}
}

// GetOutputKeys returns the output keys the chain returns.
func (c *Router) GetOutputKeys() []string {
	return []string{}
}
